#include "yaz.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <memory>
#include <string>
#include <string_view>
#include <utility>
#include <zlib.h>

#include "portable_endian.h"

namespace yaz {
ByteSourceFn MemoryByteSource(std::string_view input) {
  auto is_called = std::make_shared<bool>(false);

  return [is_called, input]() -> std::string_view {
    if (!(*is_called)) {
      *is_called = true;

      return input;
    }

    return {};
  }; // the return value
}

ByteSourceFn FileByteSource(const char *filename) {
  FILE *file = fopen(filename, "rb");
  if (file == nullptr)
    return nullptr;

  struct State {
    explicit State(FILE *f) : file(f) {}
    ~State() {}
    FILE *file;
    char buffer[65536];
  }; // in class

  auto state = std::make_shared<State>(file); // capture req. copy constructable

  return [state]() -> std::string_view {
    size_t r = fread(state->buffer, 1, sizeof(State::buffer), state->file);
    return {state->buffer, r};
  };
}

// handle header writer, talking care of little-endianness writing of fileds
namespace {
struct HeaderWriter {
  explicit HeaderWriter(char *buffer) : begin_(buffer), pos_(buffer) {}

  HeaderWriter &AddInt16(uint16_t value) {
    uint16_t val = htole16(value);
    memcpy(pos_, &val, 2);
    pos_ += 2;

    return *this;
  }

  HeaderWriter &AddInt32(uint32_t value) {
    uint32_t val = htole32(value);
    memcpy(pos_, &val, 4);
    pos_ += 4;

    return *this;
  }

  HeaderWriter &AddLiteral(std::string_view str) {
    memcpy(pos_, str.data(), str.size());
    pos_ += str.size();

    return *this;
  }

  bool Write(const ByteSinkFn &out) {
    // out: (string_view) -> bool
    // using ByteSinkFn = std::function<bool(std::string_view out)>;
    return out({begin_, static_cast<size_t>(pos_ - begin_)});
  }

  void Append(std::string *buffer) { buffer->append(begin_, pos_ - begin_); }

private:
  const char *begin_;
  char *pos_;
}; // class HeaderWriter

} // namespace

// NOTE: handle yaz.h now
struct Encoder::Impl {
  static constexpr int16_t ZIPVERSION = 20; // 2.0 pertty basic
  struct CompressResult {
    uint32_t input_crc;
    size_t input_size;
    size_t output_size;
  };

  Impl(int compression_level, ByteSinkFn out)
      : compression_level_(std::clamp(compression_level, 0, 9)),
        delegate_write_(std::move(out)), out_([this](std::string_view s) {
          output_file_offset_ += s.size();
          return delegate_write_(s);
        }) {}

  bool AddFile(std::string_view filename,
               const ByteSourceFn &content_generator) {
    if (is_finished_)
      return false; // can not add more files

    if (!content_generator)
      return false;

    ++file_count_;

    const size_t start_offset = output_file_offset_;

    const uint16_t mod_time = 0;
    const uint16_t mod_date = 0;

    bool success = HeaderWriter(scratch_space_)
                       .AddLiteral("PK\x03\x04")
                       .AddInt16(ZIPVERSION)
                       .AddInt16(0x08)
                       .AddInt16(compression_level_ == 0 ? 0 : 8)
                       .AddInt16(mod_time)
                       .AddInt16(mod_date)
                       .AddInt32(0) // CRC32: known later.
                       .AddInt32(0) // compressed size: known later.
                       .AddInt32(0) // uncompressed size: known later.
                       .AddInt16(filename.length())
                       .AddInt16(0) // extra field length
                       .AddLiteral(filename)
                       .Write(out_);

    if (!success)
      return false;

    // Data output
    const CompressResult compress_result =
        compression_level_ == 0 ? CopyDataToOutput(content_generator)
                                : CompressDataToOutput(content_generator);
    // Asemble Data Descriptor after file with known CRC and size
    success = HeaderWriter(scratch_space_)
                  .AddInt32(compress_result.input_crc)
                  .AddInt32(compress_result.output_size)
                  .AddInt32(compress_result.input_size)
                  .Write(out_);

    HeaderWriter(scratch_space_)
        .AddLiteral("PK\x01\x02")
        .AddInt16(ZIPVERSION)
        .AddInt16(ZIPVERSION)
        .AddInt16(0x08) // flag
        .AddInt16(compression_level_ == 0 ? 0 : 8)
        .AddInt16(mod_time)
        .AddInt16(mod_date) // QUES:modify time and date?
        .AddInt32(compress_result.input_crc)
        .AddInt32(compress_result.output_size)
        .AddInt32(compress_result.input_size)
        .AddInt16(filename.length())
        .AddInt16(0) // TODO: add desc here
        .AddInt16(0)
        .AddInt16(0)
        .AddInt16(0)
        .AddInt32(0)
        .AddInt32(start_offset)
        .AddLiteral(filename)
        .Append(&central_dir_data_);

    return success;
  }

  bool Finish() {
    if (is_finished_)
      return false;

    is_finished_ = true;

    const size_t start_offset = output_file_offset_;
    if (!out_({central_dir_data_.data(), central_dir_data_.size()})) {
      return false;
    }

    constexpr std::string_view comment(
        "Create with yaz(yet another zip lib(not for unzip))");

    return HeaderWriter(scratch_space_)
        .AddLiteral("PK\x05\x06")
        .AddInt16(0)
        .AddInt16(0)
        .AddInt16(file_count_)
        .AddInt16(file_count_)
        .AddInt32(central_dir_data_.size())
        .AddInt32(start_offset)
        .AddInt16(comment.length())
        .AddLiteral(comment)
        .Write(out_);
  }

  CompressResult CopyDataToOutput(const ByteSourceFn &generator) {
    uint32_t crc = 0;
    size_t processed_size = 0;
    std::string_view chunk;
    while (!(chunk = generator()).empty()) {
      crc = crc32(crc, reinterpret_cast<const uint8_t *>(chunk.data()),
                  chunk.size());
      processed_size += chunk.size();
      out_(chunk);
    }

    return {crc, processed_size, processed_size};
  }
  CompressResult CompressDataToOutput(const ByteSourceFn &generator) {
    uint32_t crc = 0;
    std::string_view chunk;
    z_stream stream;
    memset(&stream, 0x00, sizeof(stream));

    deflateInit2(&stream, compression_level_, Z_DEFLATED, -15, 9, 0);

    const size_t SCRATCH_SIZE = sizeof(scratch_space_);
    do {
      chunk = generator();
      const int flush_setting = chunk.empty() ? Z_FINISH : Z_NO_FLUSH;

      if (!chunk.empty()) {
        crc = crc32(crc, reinterpret_cast<const uint8_t *>(chunk.data()),
                    chunk.size());
      }

      stream.avail_in = chunk.size();
      stream.next_in =
          reinterpret_cast<uint8_t *>(const_cast<char *>(chunk.data()));
      do {
        stream.avail_out = SCRATCH_SIZE;
        stream.next_out = reinterpret_cast<uint8_t *>(scratch_space_);
        deflate(&stream, flush_setting);
        const size_t output_size = SCRATCH_SIZE - stream.avail_out;
        if (output_size)
          out_({scratch_space_, output_size});
      } while (stream.avail_out == 0);
      // ---- end while
    } while (!chunk.empty());

    CompressResult result = {crc, stream.total_in, stream.total_out};
    deflateEnd(&stream);

    return result;
  }

private:
  const int compression_level_;
  const ByteSinkFn delegate_write_;
  const ByteSinkFn out_;

  size_t file_count_ = 0;
  size_t output_file_offset_ = 0;
  std::string central_dir_data_;
  bool is_finished_ = false;
  char scratch_space_[1 << 20]; // to assemble headers and compression data
};                              // class Encoder::Impl

Encoder::Encoder(int compression_level, ByteSinkFn out)
    : impl_(new Impl(compression_level, std::move(out))) {}

Encoder::~Encoder() { Finish(); }

bool Encoder::AddFile(std::string_view filename,
                      const ByteSourceFn &content_generator) {
  return impl_->AddFile(filename, content_generator);
}

bool Encoder::Finish() { return impl_->Finish(); }
// NOTE: end handle yaz.h

} // namespace yaz
